import os
import torch
import argparse
import torch.nn.functional as F
from torch.utils.data import DataLoader
from src.models.modnet import MODNet
from src.utils import filter_parameters_by_type
from src.trainer import supervised_training_iter, soc_adaptation_iter
from dataset.data import DatasetNoLabel, DatasetMatting
from tensorboardX import SummaryWriter
import pickle
import logging
logging.getLogger().setLevel(logging.INFO)

# 输入参数
parser = argparse.ArgumentParser(
    'Eval on freeze MODNet'
)

parser.add_argument('--log', 
    default='./logs/tb_eval_freeze',
    help='tensorboard log path'
)

parser.add_argument('--data_path', 
    default='/home/chengk/chk/data/removebg-align/normal3',
    help='image data folder'
)

parser.add_argument('--anno_path',
    default='/home/chengk/chk/data/removebg-align/normal3_removebg',
    help='anno folder'
)

parser.add_argument('--log_image_every_step',
    default=100,
    help='as param name shows'
)

parser.add_argument('--pretrained_model', 
    default='./modnet_webcam_portrait_matting.ckpt',
    help='pretrained model path'
)

args = parser.parse_args()

tb_logger = SummaryWriter(args.log)

bs = 16         # batch size
log_image_every_step = args.log_image_every_step

device = torch.device('cuda:0')
ckpt_pretrain = args.pretrained_model
ckpt = torch.load(ckpt_pretrain, map_location='cpu')

modnet_freeze = torch.nn.DataParallel(MODNet(backbone_pretrained=False)).cuda()
modnet_freeze.load_state_dict(ckpt)
modnet_freeze.eval()

# dataset = DatasetNoLabel(root=args.data_path)
dataset = DatasetMatting(root=args.data_path, root_anno=args.anno_path, return_fn=True)
dataloader = DataLoader(dataset, bs, num_workers=4, drop_last=True, shuffle=True)     # NOTE: please finish this function

dict_loss = {}
step_scalar, step_image = 0, 0
for idx, data in enumerate(dataloader):
    log_step = (idx+1) % log_image_every_step == 0
    img, matte, trimap, fn = data
    
    # print('overlap_anchor:', overlap_anchor, '\toverlap_target:', overlap_target)
    img = img.to(device)
    matte = matte.to(device)
    trimap = trimap.to(device)
    with torch.no_grad():
        _, _, pred_matte_freeze = modnet_freeze(img, True)

    loss = 0
    loss_l1 = F.l1_loss(pred_matte_freeze, matte, reduction='none')
    for i, f in enumerate(fn):
        dict_loss[f] =  ((loss_l1[i:i+1, ...] * (trimap[i: i+1, ...] == 1).float()).sum() / (1e-6+(trimap[i: i+1, ...]==1).sum()) + \
                            (loss_l1[i:i+1, ...] * (trimap[i: i+1, ...] == 0).float()).sum() / (1e-6+(trimap[i: i+1, ...]==0).sum())).item()

        # if log_step:
        #     # Tensorboard image & matte log
        #     # tb_logger.add_images
        #     img_plot = torch.cat([(img_anchor.cpu()+1)/2, 
        #                     (img_target.cpu()+1)/2,
        #                     torch.cat([pred_matte.cpu()]*3, 1),
        #                     torch.cat([pred_matte_freeze.cpu()]*3, 1),
        #                     *plot_roi
        #                     ], 0)
        #     # assert img_plot.size(0) == 3, img_plot.size()
        #     tb_logger.add_images('anchor-target-pred-freeze-roi_anchor-roi_target', img_plot, step_image)
        #     # tb_logger.add_image('image', (image[0]+1)/2, step_image)
        #     # tb_logger.add_image('fg', pred_fg[0], step_image)
        #     # tb_logger.add_image('fg_backup', pred_backup_fg[0], step_image)
        #     step_image += 1
        # # Tensorboard loss log
        # if (idx+1) % 10 == 0:
        #     tb_logger.add_scalar('soc2_loss', loss.item(),
        #                                         step_scalar)
        #     step_scalar += 1

with open('eval_freeze.pkl', 'wb') as f:
    pickle.dump(dict_loss, f)